/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "slam_ros/offline_node.h"

#include <errno.h>
#include <string.h>
#ifndef WIN32
#include <sys/resource.h>
#endif
#include <time.h>

#include <chrono>
#include <memory>
#include <utility>

#include "absl/strings/str_split.h"
#include "slam_ros/node.h"
#include "slam_ros/playable_bag.h"
#include "slam_ros/urdf_reader.h"
#include "rclcpp/rclcpp.hpp"
#include "tf2_ros/static_transform_broadcaster.h"
#include "tf2_ros/buffer.h"
#include "tf2_ros/transform_listener.h"
#include "rosgraph_msgs/msg/clock.hpp"
#include "urdf/model.h"

DEFINE_bool(collect_metrics, false,
            "Activates the collection of runtime metrics. If activated, the "
            "metrics can be accessed via a ROS service.");
DEFINE_string(configuration_directory, "",
              "First directory in which configuration files are searched, "
              "second is always the Slam installation to allow "
              "including files from there.");
DEFINE_string(
    configuration_basenames, "",
    "Comma-separated list of basenames, i.e. not containing any "
    "directory prefix, of the configuration files for each trajectory. "
    "The first configuration file will be used for node options. "
    "If less configuration files are specified than trajectories, the "
    "first file will be used the remaining trajectories.");
DEFINE_string(
    bag_filenames, "",
    "Comma-separated list of bags to process. One bag per trajectory. "
    "Any combination of simultaneous and sequential bags is supported.");
DEFINE_string(urdf_filenames, "",
              "Comma-separated list of one or more URDF files that contain "
              "static links for the sensor configuration(s).");
DEFINE_bool(use_bag_transforms, true,
            "Whether to read, use and republish transforms from bags.");
DEFINE_string(load_state_filename, "",
              "If non-empty, filename of a .pbstream file to load, containing "
              "a saved SLAM state.");
DEFINE_bool(load_frozen_state, true,
            "Load the saved state as frozen (non-optimized) trajectories.");
DEFINE_string(save_state_filename, "",
              "Explicit name of the file to which the serialized state will be "
              "written before shutdown. If left empty, the filename will be "
              "inferred from the first bagfile's name as: "
              "<bag_filenames[0]>.pbstream");
DEFINE_double(map_resolution, 0.05, "Resolution pgm map. ");

DEFINE_double(initial_pose_x, 0.0, "Initial pose x.");
DEFINE_double(initial_pose_y, 0.0, "Initial pose y.");
DEFINE_double(initial_pose_yaw, 0.0, "Initial pose yaw.");
DEFINE_bool(cover_old_map, true, "Whether to overwrite old maps");

DEFINE_bool(keep_running, false,
            "Keep running the offline node after all messages from the bag "
            "have been processed.");
DEFINE_double(skip_seconds, 0,
              "Optional amount of seconds to skip from the beginning "
              "(i.e. when the earliest bag starts.). ");

namespace slam_ros {

constexpr char kClockTopic[] = "/clock";
constexpr char kTfStaticTopic[] = "/tf_static";
constexpr char kTfTopic[] = "/tf";
constexpr double kClockPublishFrequencySec = 1. / 30.;
constexpr int kSingleThreaded = 1;
// We publish tf messages one second earlier than other messages. Under
// the assumption of higher frequency tf this should ensure that tf can
// always interpolate.
const rclcpp::Duration kDelay = rclcpp::Duration::from_seconds(1.0);

void RunOfflineNode(const MapBuilderFactory& map_builder_factory) {
  CHECK(!FLAGS_configuration_directory.empty())
      << "-configuration_directory is missing.";
  CHECK(!FLAGS_configuration_basenames.empty())
      << "-configuration_basenames is missing.";
  CHECK(!(FLAGS_bag_filenames.empty() && FLAGS_load_state_filename.empty()))
      << "-bag_filenames and -load_state_filename cannot both be unspecified.";
  
  const std::vector<std::string> bag_filenames =
      absl::StrSplit(FLAGS_bag_filenames, ',', absl::SkipEmpty());
  
  slam_ros::NodeOptions node_options;
  const std::vector<std::string> configuration_basenames =
      absl::StrSplit(FLAGS_configuration_basenames, ',', absl::SkipEmpty());
  
  std::vector<TrajectoryOptions> bag_trajectory_options(1);
  std::tie(node_options, bag_trajectory_options.at(0)) =
      LoadOptions(FLAGS_configuration_directory, configuration_basenames.at(0));

  bool cover_old_map = true;
  if (bag_filenames.size() == 1 && !FLAGS_load_state_filename.empty()) {
    RCLCPP_INFO(rclcpp::get_logger("slam_ros"), 
                "Set initial pose to trajectory 0: %f %f %f", 
                FLAGS_initial_pose_x, FLAGS_initial_pose_y, FLAGS_initial_pose_yaw);
    
    ::slam::mapping::proto::InitialTrajectoryPose initial_trajectory_pose;
    initial_trajectory_pose.set_to_trajectory_id(0);
    const ::slam::transform::Rigid3d initial_pose =
        slam::transform::Embed3D(slam::transform::Rigid2d(
            Eigen::Vector2d(FLAGS_initial_pose_x, FLAGS_initial_pose_y), 
            FLAGS_initial_pose_yaw));
    *initial_trajectory_pose.mutable_relative_pose() =
        slam::transform::ToProto(initial_pose);
    initial_trajectory_pose.set_timestamp(slam::common::ToUniversal(
      ::slam_ros::FromRos(rclcpp::Time(0, 0, RCL_ROS_TIME))));
    *bag_trajectory_options.at(0).trajectory_builder_options
        .mutable_initial_trajectory_pose() = initial_trajectory_pose;
    cover_old_map = FLAGS_cover_old_map;
  }

  for (size_t bag_index = 1; bag_index < bag_filenames.size(); ++bag_index) {
    TrajectoryOptions current_trajectory_options;
    if (bag_index < configuration_basenames.size()) {
      std::tie(std::ignore, current_trajectory_options) = LoadOptions(
          FLAGS_configuration_directory, configuration_basenames.at(bag_index));
    } else {
      current_trajectory_options = bag_trajectory_options.at(0);
    }
    bag_trajectory_options.push_back(current_trajectory_options);
  }
  
  if (bag_filenames.size() > 0) {
    CHECK_EQ(bag_trajectory_options.size(), bag_filenames.size());
  }

  // Since we preload the transform buffer, we should never have to wait for a
  // transform. When we finish processing the bag, we will simply drop any
  // remaining sensor data that cannot be transformed due to missing transforms.
  node_options.lookup_transform_timeout_sec = 0.;

  auto map_builder = map_builder_factory(node_options.map_builder_options);

  const std::chrono::time_point<std::chrono::steady_clock> start_time =
      std::chrono::steady_clock::now();

  auto node = std::make_shared<rclcpp::Node>("slam_offline_node");
  auto tf_buffer = std::make_shared<tf2_ros::Buffer>(node->get_clock());
  auto tf_listener = std::make_shared<tf2_ros::TransformListener>(*tf_buffer);

  std::vector<geometry_msgs::msg::TransformStamped> urdf_transforms;
  const std::vector<std::string> urdf_filenames =
      absl::StrSplit(FLAGS_urdf_filenames, ',', absl::SkipEmpty());
  for (const auto& urdf_filename : urdf_filenames) {
    const auto current_urdf_transforms =
        ReadStaticTransformsFromUrdf(urdf_filename, tf_buffer.get());
    urdf_transforms.insert(urdf_transforms.end(),
                           current_urdf_transforms.begin(),
                           current_urdf_transforms.end());
  }

  auto static_tf_broadcaster = std::make_shared<tf2_ros::StaticTransformBroadcaster>(node);

  if (!urdf_transforms.empty()) {
    static_tf_broadcaster->sendTransform(urdf_transforms);
  }

  auto clock_publisher = node->create_publisher<rosgraph_msgs::msg::Clock>(
      kClockTopic, rclcpp::QoS(1));

  auto clock_republish_timer = node->create_wall_timer(
      std::chrono::duration<double>(kClockPublishFrequencySec),
      [clock_publisher]() {
        auto clock_msg = std::make_unique<rosgraph_msgs::msg::Clock>();
        clock_msg->clock = rclcpp::Time(0);
        clock_publisher->publish(std::move(clock_msg));
      });
  clock_republish_timer->cancel();

  slam_ros::Node slam_node(node_options, std::move(map_builder), tf_buffer.get(),
                          FLAGS_collect_metrics);
  
  if (!FLAGS_load_state_filename.empty()) {
    slam_node.LoadState(FLAGS_load_state_filename, FLAGS_load_frozen_state);
  }

  auto tf_publisher = node->create_publisher<tf2_msgs::msg::TFMessage>(
      kTfTopic, rclcpp::QoS(1));

  std::vector<
      std::set<slam::mapping::TrajectoryBuilderInterface::SensorId>>
      bag_expected_sensor_ids;
  if (configuration_basenames.size() == 1) {
    const auto current_bag_expected_sensor_ids =
        slam_node.ComputeDefaultSensorIdsForMultipleBags(
            {bag_trajectory_options.front()});
    bag_expected_sensor_ids = {bag_filenames.size(),
                              current_bag_expected_sensor_ids.front()};
  } else {
    bag_expected_sensor_ids =
        slam_node.ComputeDefaultSensorIdsForMultipleBags(bag_trajectory_options);
  }
  CHECK_EQ(bag_expected_sensor_ids.size(), bag_filenames.size());

  std::map<std::pair<int /* bag_index */, std::string>,
           slam::mapping::TrajectoryBuilderInterface::SensorId>
      bag_topic_to_sensor_id;
  PlayableBagMultiplexer playable_bag_multiplexer;
  
  for (size_t current_bag_index = 0; current_bag_index < bag_filenames.size();
       ++current_bag_index) {
    const std::string& bag_filename = bag_filenames.at(current_bag_index);
    if (!rclcpp::ok()) {
      return;
    }
    
    for (const auto& expected_sensor_id :
         bag_expected_sensor_ids.at(current_bag_index)) {
      const auto bag_resolved_topic = std::make_pair(
          static_cast<int>(current_bag_index),
          std::string(node->get_name()) + "/" + expected_sensor_id.id);
      if (bag_topic_to_sensor_id.count(bag_resolved_topic) != 0) {
        RCLCPP_ERROR(node->get_logger(), 
                    "Sensor %s of bag %ld resolves to topic %s which is already used by sensor %s",
                    expected_sensor_id.id.c_str(), current_bag_index,
                    bag_resolved_topic.second.c_str(),
                    bag_topic_to_sensor_id.at(bag_resolved_topic).id.c_str());
      }
      bag_topic_to_sensor_id[bag_resolved_topic] = expected_sensor_id;
    }

    playable_bag_multiplexer.AddPlayableBag(PlayableBag(
        bag_filename, current_bag_index, rclcpp::Time(0), rclcpp::Time(0), kDelay,
        [tf_publisher, tf_buffer](const std::shared_ptr<rosbag2_storage::SerializedBagMessage>& msg) {
          if (msg->topic_name == kTfTopic || msg->topic_name == kTfStaticTopic) {
            if (FLAGS_use_bag_transforms) {
              auto tf_message = std::make_shared<tf2_msgs::msg::TFMessage>();
              tf_message->transforms.push_back(geometry_msgs::msg::TransformStamped());
              tf_publisher->publish(*tf_message);

              try {
                tf_buffer->setTransform(tf_message->transforms.front(), "unused_authority",
                                      msg->topic_name == kTfStaticTopic);
              } catch (const tf2::TransformException& ex) {
                RCLCPP_WARN(rclcpp::get_logger("slam_ros"), "%s", ex.what());
              }
            }
            return false;
          } else {
            return true;
          }
        }));
  }

  std::set<std::string> bag_topics;
  std::stringstream bag_topics_string;
  for (const auto& topic : playable_bag_multiplexer.topics()) {
    std::string resolved_topic = std::string(node->get_name()) + "/" + topic;
    bag_topics.insert(resolved_topic);
    bag_topics_string << resolved_topic << ",";
  }
  
  bool print_topics = false;
  for (const auto& entry : bag_topic_to_sensor_id) {
    const std::string& resolved_topic = entry.first.second;
    if (bag_topics.count(resolved_topic) == 0) {
      RCLCPP_WARN(node->get_logger(), 
                 "Expected resolved topic \"%s\" not found in bag file(s).",
                 resolved_topic.c_str());
      print_topics = true;
    }
  }
  
  if (print_topics) {
    RCLCPP_WARN(node->get_logger(), 
               "Available topics in bag file(s) are %s",
               bag_topics_string.str().c_str());
  }

  std::unordered_map<int, int> bag_index_to_trajectory_id;
  const rclcpp::Time begin_time =
      playable_bag_multiplexer.IsMessageAvailable()
          ? playable_bag_multiplexer.PeekMessageTime()
          : rclcpp::Time(0);
  
  rclcpp::executors::SingleThreadedExecutor executor;
  executor.add_node(node);
  
  while (playable_bag_multiplexer.IsMessageAvailable()) {
    if (!rclcpp::ok()) {
      return;
    }

    const auto next_msg_tuple = playable_bag_multiplexer.GetNextMessage();
    const auto& msg = std::get<0>(next_msg_tuple);
    const int bag_index = std::get<1>(next_msg_tuple);
    const bool is_last_message_in_bag = std::get<2>(next_msg_tuple);

    if (msg->time_stamp < 
        (begin_time + rclcpp::Duration::from_seconds(FLAGS_skip_seconds)).nanoseconds()) {
      continue;
    }

    int trajectory_id;
    if (bag_index_to_trajectory_id.count(bag_index) == 0) {
      trajectory_id =
          slam_node.AddOfflineTrajectory(bag_expected_sensor_ids.at(bag_index),
                                        bag_trajectory_options.at(bag_index));
      CHECK(bag_index_to_trajectory_id
                .emplace(std::piecewise_construct,
                         std::forward_as_tuple(bag_index),
                         std::forward_as_tuple(trajectory_id))
                .second);
      RCLCPP_INFO(node->get_logger(), 
                 "Assigned trajectory %d to bag %s",
                 trajectory_id, bag_filenames.at(bag_index).c_str());
    } else {
      trajectory_id = bag_index_to_trajectory_id.at(bag_index);
    }

    const auto bag_topic = std::make_pair(
        bag_index,
        std::string(node->get_name()) + "/" + msg->topic_name);
    auto it = bag_topic_to_sensor_id.find(bag_topic);
    if (it != bag_topic_to_sensor_id.end()) {
      const std::string& sensor_id = it->second.id;
      if (msg->topic_name == "sensor_msgs/msg/LaserScan") {
        auto scan = std::make_shared<sensor_msgs::msg::LaserScan>();
        slam_node.HandleLaserScanMessage(trajectory_id, sensor_id, scan);
      }
      if (msg->topic_name == "sensor_msgs/msg/MultiEchoLaserScan") {
        auto scan = std::make_shared<sensor_msgs::msg::MultiEchoLaserScan>();
        slam_node.HandleMultiEchoLaserScanMessage(trajectory_id, sensor_id, scan);
      }
      if (msg->topic_name == "sensor_msgs/msg/PointCloud2") {
        auto cloud = std::make_shared<sensor_msgs::msg::PointCloud2>();
        slam_node.HandlePointCloud2Message(trajectory_id, sensor_id, cloud);
      }
      if (msg->topic_name == "sensor_msgs/msg/Imu") {
        auto imu = std::make_shared<sensor_msgs::msg::Imu>();
        slam_node.HandleImuMessage(trajectory_id, sensor_id, imu);
      }
      if (msg->topic_name == "agv_msgs/msg/OdometryLite") {
        auto odom = std::make_shared<agv_msgs::msg::OdometryLite>();
        slam_node.HandleOdometryMessage(trajectory_id, sensor_id, odom);
      }
      if (msg->topic_name == "sensor_msgs/msg/NavSatFix") {
        auto navsat = std::make_shared<sensor_msgs::msg::NavSatFix>();
        slam_node.HandleNavSatFixMessage(trajectory_id, sensor_id, navsat);
      }
      if (msg->topic_name == "slam_ros_msgs/msg/LandmarkList") {
        auto landmarks = std::make_shared<slam_ros_msgs::msg::LandmarkList>();
        slam_node.HandleLandmarkMessage(trajectory_id, sensor_id, landmarks);
      }
    }
    
    auto clock_msg = std::make_unique<rosgraph_msgs::msg::Clock>();
    clock_msg->clock = rclcpp::Time(msg->time_stamp);
    clock_publisher->publish(std::move(clock_msg));

    if (is_last_message_in_bag) {
      slam_node.FinishTrajectory(trajectory_id);
    }
    
    executor.spin_some();
  }

  clock_republish_timer->reset();
  slam_node.RunFinalOptimization();

  const std::chrono::time_point<std::chrono::steady_clock> end_time =
      std::chrono::steady_clock::now();
  const double wall_clock_seconds =
      std::chrono::duration_cast<std::chrono::duration<double>>(end_time -
                                                                start_time)
          .count();

  RCLCPP_INFO(node->get_logger(), 
             "Elapsed wall clock time: %f s", wall_clock_seconds);
#ifdef __linux__
  timespec cpu_timespec = {};
  clock_gettime(CLOCK_PROCESS_CPUTIME_ID, &cpu_timespec);
  RCLCPP_INFO(node->get_logger(), 
             "Elapsed CPU time: %f s", 
             (cpu_timespec.tv_sec + 1e-9 * cpu_timespec.tv_nsec));
  rusage usage;
  CHECK_EQ(getrusage(RUSAGE_SELF, &usage), 0) << strerror(errno);
  RCLCPP_INFO(node->get_logger(), 
             "Peak memory usage: %ld KiB", usage.ru_maxrss);
#endif

  if (rclcpp::ok() &&
      !(bag_filenames.empty() && FLAGS_save_state_filename.empty())) {
    const std::string map_state_save_file = 
        FLAGS_save_state_filename + ".pbstream";
    const std::string state_output_filename =
        FLAGS_save_state_filename.empty()
            ? absl::StrCat(bag_filenames.front(), ".pbstream")
            : map_state_save_file;
    const std::string pgm_filename =
        FLAGS_save_state_filename.empty()
            ? bag_filenames.front()
            : FLAGS_save_state_filename;
    RCLCPP_INFO(node->get_logger(), 
               "Writing state to '%s'...", state_output_filename.c_str());
    slam_node.SerializeState(state_output_filename, cover_old_map,
                            true /* include_unfinished_submaps */);
  }
  
  if (FLAGS_keep_running) {
    RCLCPP_INFO(node->get_logger(), 
               "Finished processing and waiting for shutdown.");
    executor.spin();
  }
}

}  // namespace slam_ros